import torch
import torch.nn.functional as F

# 定义一个5x5的输入矩阵（模拟单通道图像）
input = torch.tensor([[1, 2, 0, 3, 1],
                      [0, 1, 2, 3, 1],
                      [1, 2, 1, 0, 0],
                      [5, 2, 3, 1, 1],
                      [2, 1, 0, 1, 1]
                      ])

# 定义一个3x3的卷积核（滤波器）
kernel = torch.tensor([[1, 2, 1],
                       [0, 1, 0],
                       [2, 1, 0]])

# 将输入和卷积核reshape为4D张量（batch_size=1, channels=1, height, width）
input = torch.reshape(input, (1, 1, 5, 5))  # 形状变为 torch.Size([1, 1, 5, 5])
kernel = torch.reshape(kernel, (1, 1, 3, 3))  # 形状变为 torch.Size([1, 1, 3, 3])

# 打印输入和卷积核
print(f"输入input：{input}")
print(f"核kernel：{kernel}")
print("\n")

# 情况1：stride=1, padding=0（默认参数）
output = F.conv2d(input, kernel, stride=1)
print(f"stride=1,padding=0,输出output：{output}")
print(f"stride=1,padding=0,输出的形状：{output.shape}")  # 输出尺寸公式：(5-3)/1+1=3
print("\n")

# 情况2：stride=2, padding=0（步长为2）
output2 = F.conv2d(input, kernel, stride=2)
print(f"stride=2,padding=0,输出output2：{output2}")
print(f"stride=2,padding=0,输出的形状：{output.shape}")  # 输出尺寸公式：(5-3)/2+1=2
print("\n")

# 情况3：stride=1, padding=1（边缘填充1圈0）
output3 = F.conv2d(input, kernel, stride=1, padding=1)
print(f"stride=1,padding=1,输出output3：{output3}")
print(f"stride=1,padding=1,输出的形状：{output3.shape}")  # 输出尺寸公式：(5+2-3)/1+1=5
print("\n")